package com.pshdhx.aiagent.rag;

import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.ai.rag.Query;
import org.springframework.ai.rag.preretrieval.query.transformation.QueryTransformer;
import org.springframework.ai.rag.preretrieval.query.transformation.RewriteQueryTransformer;
import org.springframework.context.annotation.Configuration;

/**
 * @Author pansd
 * @Date 2025-07-08 9:32
 * @Des 提示词重写器
 */
@Configuration
public class QueryReWriteConfig {

    private final QueryTransformer queryTransformer;

    public QueryReWriteConfig(ChatModel dashscopeChatModel) {
        ChatClient.Builder builder = ChatClient.builder(dashscopeChatModel);
        queryTransformer = RewriteQueryTransformer.builder().chatClientBuilder(builder).build();
    }

    public String doRewriteQuery(String prompt){
        Query reWritedQuery = this.queryTransformer.transform(new Query(prompt));
        return reWritedQuery.text();
    }
}
